package com.yx.vectordb.example;

import com.yx.vectordb.api.IVectorDB;
import com.yx.vectordb.impl.VectorSearchResult;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.allminilml6v2.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.model.output.Response;

import java.io.File;
import java.util.List;

/**
 * 向量数据库使用示例
 * 演示了如何使用向量数据库进行文本向量化存储和相似度搜索
 *
 * @author Dingwei He
 */
public class VectorDBExample {
    /** 使用 AllMiniLmL6V2 模型进行文本向量化 */
    static EmbeddingModel embeddingModel = new AllMiniLmL6V2EmbeddingModel();

    /**
     * 示例程序入口
     * 展示了向量数据库的基本操作，包括：
     * 1. 初始化数据库
     * 2. 添加向量数据
     * 3. 执行相似度搜索
     */
    public static void main(String[] args) throws Exception {
        try {
            // 初始化向量数据库，指定存储路径
            File dbFile = new File("./data/test_384");
            IVectorDB<SimpleMeta> db = IVectorDB.open(dbFile, SimpleMeta.class, embeddingModel.dimension());

            // 示例1：添加第一条向量数据
            SimpleMeta simpleMeta = new SimpleMeta();
            simpleMeta.setDocument("今天的天气怎么样？适合出去走走吗？");
            simpleMeta.setUserName("Lucy");
            db.addVector("1", embedding(simpleMeta.getDocument()), simpleMeta);

            // 示例2：添加第二条向量数据
            SimpleMeta simpleMeta1 = new SimpleMeta();
            simpleMeta1.setDocument("今天我想去打球，不知道有没有人想去打");
            simpleMeta1.setUserName("Lily");
            db.addVector("2", embedding(simpleMeta1.getDocument()), simpleMeta1);

            // 示例3：执行向量相似度搜索（不带过滤条件）
            {
                List<VectorSearchResult<SimpleMeta>> list = db.search(
                    embedding("天气情况"),  // 搜索文本
                    0.7f,                  // 相似度阈值
                    10,                    // 返回结果数量
                    null                   // 过滤条件
                );
                for (VectorSearchResult vectorSearchResult : list) {
                    System.out.println(vectorSearchResult);
                }
            }

            // 示例4：执行向量相似度搜索（带过滤条件）
            {
                List<VectorSearchResult<SimpleMeta>> list = db.search(
                    embedding("天气情况"),  // 搜索文本
                    0.7f,                  // 相似度阈值
                    10,                    // 返回结果数量
                    " userName like 'h%' or document like '%api%' "             // 过滤条件, 这里就是sql的where语句 作用在meta语句上
                );
                for (VectorSearchResult vectorSearchResult : list) {
                    System.out.println(vectorSearchResult);
                }
            }

        } catch (Exception ex) {
            // 异常处理
        } finally {
            // 资源清理
        }
    }

    /**
     * 将文本转换为向量
     * 使用 AllMiniLmL6V2 模型将输入文本转换为向量表示
     *
     * @param input 输入文本
     * @return 文本的向量表示
     */
    public static float[] embedding(String input) {
        Response<Embedding> response = embeddingModel.embed(input);
        Embedding embedding = response.content();
        return embedding.vector();
    }
}